// #include <rkllm_tools.hpp>
#include "../inc/rkllm_tools.hpp"
#include <iostream>
#include <thread>

const std::string R1_PROMPT_TEXT_PREFIX = "Aconversation between User and Assistant. The user asks a question, and the Assistant solves it."
                                          "The assistant first thinks about the reasoning process in the mind and then provides the user"
                                          "with the answer. The reasoning process and answer are enclosed within <think> </think> and"
                                          " <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think>"
                                          " <answer> answer here </answer>. ";

const std::string R1_PROMPT_TEXT_POSTFIX = "Assistant: ";


const std::string PROMPT_TEXT_PREFIX = "<|im_start|>system You are a helpful assistant. <|im_end|> <|im_start|>";
const std::string PROMPT_TEXT_POSTFIX = "<|im_end|><|im_start|>assistant";

int main(int argc, const char* argv[]) {
  if (argc == 2) {
    printf("%s:%d\n", __FILE__, __LINE__);
    std::filesystem::path path = argv[1];
    rkllmt::RKLLMInfer::Param param;
    std::string prefix;
    std::string suffix;
    if (path.filename().string().find("R1") != std::string::npos) {
      param.temperature = 0.6;
      param.max_context_len = 16000;
      param.max_new_tokens  = 16000;
      prefix = R1_PROMPT_TEXT_PREFIX;
      suffix = R1_PROMPT_TEXT_POSTFIX;
    } else {
      prefix = PROMPT_TEXT_PREFIX;
      suffix = PROMPT_TEXT_POSTFIX;
    }
    rkllmt::RKLLMInfer infer(path, param);
    std::string global_text = "";
    while (true) {
      std::string input_str;
      printf("\n");
      printf("user: ");
      std::getline(std::cin, input_str);
      if (input_str == "exit") break;
      if (input_str == "new") {
        global_text = "";
        continue;
      }
      global_text += "user: " + input_str + "\n";
      std::shared_ptr<rkllmt::AtomicUserData> user_data = nullptr;
      infer.RunAsync(user_data, prefix + global_text + suffix);
      std::cout << "===============================\n";
      std::cout << prefix + global_text + suffix << "\n";
      std::cout << "===============================\n";
      global_text += suffix;
      printf("robot: ");
      std::string text;
      int total_token_count = 0;
      int total_get_count = 0;
      for (int state = 0; state == 0; total_token_count += user_data->Pop(text, state)) {
        std::cout << text;
        global_text += text;
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
        ++total_get_count;
      }
      global_text += "\n";
      float token_per_second = static_cast<float>(total_token_count) / (total_get_count * 0.1f);
      printf("\n%.2f token/s", token_per_second);
    }
  }
  
}
